#
#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

import logging
import json
import uuid

import valkey as redis
from rag import settings
from rag.utils import singleton
from valkey.lock import Lock
# trio类似asyncio的异步框架
import trio

# 定义 RedisMsg 类，用于封装从 Redis Stream 中读取的消息，并提供确认和获取消息内容的方法
class RedisMsg:
    # 构造函数接收消费者对象、队列名、消费者组名、消息ID和原始消息内容，初始化 RedisMsg 实例
    def __init__(self, consumer, queue_name, group_name, msg_id, message):
        # 保存 Redis 消费者对象用于后续确认操作
        self.__consumer = consumer
        # 保存当前消息所属的队列名称
        self.__queue_name = queue_name
        # 保存消费者组名称，用于消息确认时使用
        self.__group_name = group_name
        # 保存当前消息的唯一标识符（msg_id）
        self.__msg_id = msg_id
        # 解析并保存消息体内容为 JSON 对象
        self.__message = json.loads(message["message"])

    # 确认该条消息已被成功处理，防止其再次被消费
    def ack(self):
        try:
            # 调用 Redis xack 方法对当前消息进行确认
            self.__consumer.xack(self.__queue_name, self.__group_name, self.__msg_id)
            # 返回确认成功标志
            return True
        except Exception as e:
            # 记录确认失败的日志信息，包含队列名和异常详情
            logging.warning("[EXCEPTION]ack" + str(self.__queue_name) + "||" + str(e))
        # 返回确认失败标志
        return False

    # 获取解析后的消息内容（字典格式）
    def get_message(self):
        return self.__message

    # 获取当前消息的 ID
    def get_msg_id(self):
        return self.__msg_id


# 使用单例模式装饰器，确保RedisDB在整个应用中只有一个实例
@singleton
class RedisDB:
    # 用于存储注册的Lua脚本对象（删除键时判断值是否相等）
    lua_delete_if_equal = None
    # Lua脚本定义：当指定key对应的值等于给定值时，删除该key
    LUA_DELETE_IF_EQUAL_SCRIPT = """
        local current_value = redis.call('get', KEYS[1])
        if current_value and current_value == ARGV[1] then
            redis.call('del', KEYS[1])
            return 1
        end
        return 0
    """

    # 初始化方法，初始化Redis连接和配置
    def __init__(self):
        # 存储Redis客户端连接
        self.REDIS = None
        # 从全局设置中读取Redis配置
        self.config = settings.REDIS
        # 打开或重新建立Redis连接
        self.__open__()

    # 注册自定义的Lua脚本到Redis客户端
    def register_scripts(self) -> None:
        cls = self.__class__
        client = self.REDIS
        cls.lua_delete_if_equal = client.register_script(cls.LUA_DELETE_IF_EQUAL_SCRIPT)

    # 建立Redis连接并注册脚本
    def __open__(self):
        try:
            # 创建一个StrictRedis实例，支持连接Redis服务器
            self.REDIS = redis.StrictRedis(
                # 绑定host
                host=self.config["host"].split(":")[0],
                # 绑定ip
                port=int(self.config.get("host", ":6379").split(":")[1]),
                # 数据库名
                db=int(self.config.get("db", 1)),
                # 密码
                password=self.config.get("password"),
                decode_responses=True,
            )
            # 注册所有自定义的Lua脚本
            self.register_scripts()
        except Exception:
            # 如果连接失败，记录警告日志
            logging.warning("Redis can't be connected.")
        return self.REDIS

    # 检查Redis健康状态
    def health(self):
        # 发送ping命令测试连接是否正常
        self.REDIS.ping()
        a, b = "xx", "yy"
        # 设置一个带过期时间的键值对
        self.REDIS.set(a, b, 3)

        # 验证写入和读取操作是否正常
        if self.REDIS.get(a) == b:
            return True

    # 判断当前Redis连接是否存在
    def is_alive(self):
        return self.REDIS is not None

    # 检查指定的key是否存在
    def exist(self, k):
        if not self.REDIS:
            return
        try:
            return self.REDIS.exists(k)
        except Exception as e:
            # 日志记录异常信息，并尝试重新连接Redis
            logging.warning("RedisDB.exist " + str(k) + " got exception: " + str(e))
            self.__open__()

    # 获取指定key的值
    def get(self, k):
        if not self.REDIS:
            return
        try:
            return self.REDIS.get(k)
        except Exception as e:
            # 记录获取失败的日志，并尝试重新连接
            logging.warning("RedisDB.get " + str(k) + " got exception: " + str(e))
            self.__open__()

    # 将Python对象序列化后存入Redis
    def set_obj(self, k, obj, exp=3600):
        try:
            # 使用json将对象序列化为字符串后存储，并设置过期时间
            self.REDIS.set(k, json.dumps(obj, ensure_ascii=False), exp)
            return True
        except Exception as e:
            # 异常处理并尝试重连
            logging.warning("RedisDB.set_obj " + str(k) + " got exception: " + str(e))
            self.__open__()
        return False

    # 设置字符串类型的键值对，并设置过期时间
    def set(self, k, v, exp=3600):
        try:
            self.REDIS.set(k, v, exp)
            return True
        except Exception as e:
            # 异常时记录日志并尝试重新连接
            logging.warning("RedisDB.set " + str(k) + " got exception: " + str(e))
            self.__open__()
        return False

    # 向集合类型(key)中添加成员(member)
    def sadd(self, key: str, member: str):
        try:
            self.REDIS.sadd(key, member)
            return True
        except Exception as e:
            # 出现错误时记录日志，并尝试重新连接
            logging.warning("RedisDB.sadd " + str(key) + " got exception: " + str(e))
            self.__open__()
        return False

    # Srem 命令用于移除集合中的一个或多个成员元素，不存在的成员元素会被忽略。
    def srem(self, key: str, member: str):
        try:
            self.REDIS.srem(key, member)
            return True
        except Exception as e:
            logging.warning("RedisDB.srem " + str(key) + " got exception: " + str(e))
            self.__open__()
        return False

    # 获取集合(key)中的所有成员 Redis Smembers 命令返回集合中的所有的成员。 不存在的集合 key 被视为空集合
    def smembers(self, key: str):
        try:
            res = self.REDIS.smembers(key)
            return res
        except Exception as e:
            logging.warning(
                "RedisDB.smembers " + str(key) + " got exception: " + str(e)
            )
            self.__open__()
        return None

    # 向有序集合(key)中添加一个成员及其对应的分数(score)
    def zadd(self, key: str, member: str, score: float):
        try:
            self.REDIS.zadd(key, {member: score})
            return True
        except Exception as e:
            logging.warning("RedisDB.zadd " + str(key) + " got exception: " + str(e))
            self.__open__()
        return False

    # 统计有序集合(key)中分数范围在[min, max]之间的成员数量
    def zcount(self, key: str, min: float, max: float):
        try:
            res = self.REDIS.zcount(key, min, max)
            return res
        except Exception as e:
            logging.warning("RedisDB.zcount " + str(key) + " got exception: " + str(e))
            self.__open__()
        return 0

    # 从有序集合(key)中弹出count个最小分数的成员
    def zpopmin(self, key: str, count: int):
        try:
            res = self.REDIS.zpopmin(key, count)
            return res
        except Exception as e:
            logging.warning("RedisDB.zpopmin " + str(key) + " got exception: " + str(e))
            self.__open__()
        return None

    # 获取有序集合中指定分数范围内的所有成员（升序）
    def zrangebyscore(self, key: str, min: float, max: float):
        try:
            res = self.REDIS.zrangebyscore(key, min, max)
            return res
        except Exception as e:
            logging.warning(
                "RedisDB.zrangebyscore " + str(key) + " got exception: " + str(e)
            )
            self.__open__()
        return None

    # 使用事务机制设置一个键值对，并保证原子性（仅当键不存在时设置成功）
    def transaction(self, key, value, exp=3600):
        try:
            pipeline = self.REDIS.pipeline(transaction=True)
            pipeline.set(key, value, exp, nx=True)
            pipeline.execute()
            return True
        except Exception as e:
            logging.warning(
                "RedisDB.transaction " + str(key) + " got exception: " + str(e)
            )
            self.__open__()
        return False

    # 向流（stream）类型队列中添加消息，尝试最多三次直到成功
    def queue_product(self, queue, message) -> bool:
        for _ in range(3):
            try:
                payload = {"message": json.dumps(message)}
                self.REDIS.xadd(queue, payload)
                return True
            except Exception as e:
                logging.exception(
                    "RedisDB.queue_product " + str(queue) + " got exception: " + str(e)
                )
        return False

    # 作为消费者从指定的流队列中读取一条消息，并自动创建消费者组（如果不存在）
    def queue_consumer(self, queue_name, group_name, consumer_name, msg_id=b">") -> RedisMsg:
        """https://redis.io/docs/latest/commands/xreadgroup/"""
        try:
            # 获取该流的消费者组信息
            group_info = self.REDIS.xinfo_groups(queue_name)
            # 如果目标消费者组不存在，则创建它
            if not any(gi["name"] == group_name for gi in group_info):
                self.REDIS.xgroup_create(queue_name, group_name, id="0", mkstream=True)
            # 构建xreadgroup命令参数
            args = {
                "groupname": group_name,
                "consumername": consumer_name,
                "count": 1,
                "block": 5,
                "streams": {queue_name: msg_id},
            }
            # 执行读取消息的操作
            messages = self.REDIS.xreadgroup(**args)
            # 如果没有消息返回None
            if not messages:
                return None
            stream, element_list = messages[0]
            # 如果元素列表为空也返回None
            if not element_list:
                return None
            msg_id, payload = element_list[0]
            # 返回封装好的消息对象
            res = RedisMsg(self.REDIS, queue_name, group_name, msg_id, payload)
            return res
        except Exception as e:
            # 捕获异常并记录日志，'no such key'异常被忽略
            if str(e) == 'no such key':
                pass
            else:
                logging.exception(
                    "RedisDB.queue_consumer "
                    + str(queue_name)
                    + " got exception: "
                    + str(e)
                )
        return None

    # 迭代获取多个队列中未确认的消息
    def get_unacked_iterator(self, queue_names: list[str], group_name, consumer_name):
        try:
            # 遍历所有传入的队列名称
            for queue_name in queue_names:
                try:
                    # 获取当前队列的消费者组信息
                    group_info = self.REDIS.xinfo_groups(queue_name)
                except Exception as e:
                    # 如果队列不存在则跳过
                    if str(e) == 'no such key':
                        logging.warning(f"RedisDB.get_unacked_iterator queue {queue_name} doesn't exist")
                        continue
                # 如果消费者组不存在于当前队列中，则跳过该队列
                if not any(gi["name"] == group_name for gi in group_info):
                    logging.warning(f"RedisDB.get_unacked_iterator queue {queue_name} group {group_name} doesn't exist")
                    continue
                current_min = 0
                while True:
                    # 使用队列消费者方法获取下一条未确认消息
                    payload = self.queue_consumer(queue_name, group_name, consumer_name, current_min)
                    if not payload:
                        break
                    current_min = payload.get_msg_id()
                    # 记录已消费的消息ID
                    logging.info(f"RedisDB.get_unacked_iterator {queue_name} {consumer_name} {current_min}")
                    # 通过迭代器返回结果
                    yield payload
        except Exception:
            # 捕获全局异常并重新连接Redis
            logging.exception(
                "RedisDB.get_unacked_iterator got exception: "
            )
            self.__open__()

    # 获取某个队列中处于pending状态的消息
    # 获取某个队列中处于 pending 状态的消息列表（即被消费者读取但尚未确认的消息）
    def get_pending_msg(self, queue, group_name):
        try:
            # 查询pending范围内的消息
            # 使用 Redis 的 XPENDING 命令获取指定队列和消费者组中的 pending 消息，范围是全部（'-' 到 '+'），最多取10条
            messages = self.REDIS.xpending_range(queue, group_name, '-', '+', 10)
            return messages
        except Exception as e:
            # 忽略"No such key"异常，其他异常记录日志
            # 捕获所有异常，如果异常信息不是 "No such key"，则记录警告日志
            if 'No such key' not in (str(e) or ''):
                # 记录 Redis 异常信息，便于排查问题
                logging.warning(
                    "RedisDB.get_pending_msg " + str(queue) + " got exception: " + str(e)
                )
        # 如果发生异常（如队列不存在），返回空列表以保证函数健壮性
        return []

    # 将某条特定ID的消息重新放回队列中，并标记为已ACK
    def requeue_msg(self, queue: str, group_name: str, msg_id: str):
        try:
            # 获取指定ID的消息内容
            messages = self.REDIS.xrange(queue, msg_id, msg_id)
            if messages:
                # 重新添加消息到队列尾部
                self.REDIS.xadd(queue, messages[0][1])
                # 标记原消息为已确认
                self.REDIS.xack(queue, group_name, msg_id)
        except Exception as e:
            # 异常处理并记录日志
            logging.warning(
                "RedisDB.get_pending_msg " + str(queue) + " got exception: " + str(e)
            )

    # 获取指定队列的消费者组信息
    def queue_info(self, queue, group_name) -> dict | None:
        try:
            # 获取该队列的所有消费者组信息
            groups = self.REDIS.xinfo_groups(queue)
            for group in groups:
                # 查找匹配的消费者组
                if group["name"] == group_name:
                    return group
        except Exception as e:
            # 捕获异常并记录日志
            logging.warning(
                "RedisDB.queue_info " + str(queue) + " got exception: " + str(e)
            )
        return None

    # 原子性地判断key对应的值是否等于期望值，如果是则删除该key
    def delete_if_equal(self, key: str, expected_value: str) -> bool:
        """
        Do follwing atomically:
        Delete a key if its value is equals to the given one, do nothing otherwise.
        """
        return bool(self.lua_delete_if_equal(keys=[key], args=[expected_value], client=self.REDIS))

    # 删除指定的Redis键
    def delete(self, key) -> bool:
        try:
            # 调用Redis的delete命令
            self.REDIS.delete(key)
            return True
        except Exception as e:
            # 异常时记录日志并尝试重新连接Redis
            logging.warning("RedisDB.delete " + str(key) + " got exception: " + str(e))
            self.__open__()
        return False
    
    
# 实例化
REDIS_CONN = RedisDB()


# 定义Redis分布式锁类，用于实现跨进程/节点的互斥锁机制
class RedisDistributedLock:
    # 初始化锁的参数：锁的key、值、超时时间、阻塞等待时间
    def __init__(self, lock_key, lock_value=None, timeout=10, blocking_timeout=1):
        # 保存锁的唯一标识key
        self.lock_key = lock_key
        # 如果未提供锁值，则生成一个随机UUID作为锁的值
        if lock_value:
            self.lock_value = lock_value
        else:
            self.lock_value = str(uuid.uuid4())
        # 设置锁的过期时间（秒）
        self.timeout = timeout
        # 创建底层的Redis Lock对象，用于实际加锁和解锁操作
        self.lock = Lock(REDIS_CONN.REDIS, lock_key, timeout=timeout, blocking_timeout=blocking_timeout)

    # 获取锁的方法，尝试获取锁，并在获取前删除可能存在的旧锁标记
    def acquire(self):
        # 尝试删除当前锁key，仅当其值等于预期值（避免误删他人的锁）
        REDIS_CONN.delete_if_equal(self.lock_key, self.lock_value)
        # 调用底层Lock对象尝试获取锁，并传入token（锁值）
        return self.lock.acquire(token=self.lock_value)

    # 异步自旋获取锁的方法，持续尝试直到成功获取锁
    async def spin_acquire(self):
        # 在尝试获取锁之前，先尝试删除可能残留的锁标记
        REDIS_CONN.delete_if_equal(self.lock_key, self.lock_value)
        while True:
            # 尝试获取锁，如果成功则退出循环
            if self.lock.acquire(token=self.lock_value):
                break
            # 如果未获取到锁，则异步休眠10秒后重试
            await trio.sleep(10)

    # 释放锁的方法
    def release(self):
        # 调用delete_if_equal方法，只有当前锁的值匹配时才会删除锁key，实现安全释放
        REDIS_CONN.delete_if_equal(self.lock_key, self.lock_value)
